from transformers import pipeline


class ModelManager:
    """
    模型管理器类
    用于加载和管理Hugging Face的transformer模型
    """

    def __init__(self, config):
        """
        初始化模型管理器
        参数:
            config (dict): 配置字典，包含以下键:
                - task_type: 任务类型(如'text-classification', 'text-generation'等)
                - model_name: 要加载的模型名称
        """
        self.pipeline = pipeline(task=config["task_type"], model=config["model_name"])

    def predict(self, text):
        """
        使用加载的模型进行预测
        参数:
            text (str): 需要进行预测的输入文本
        返回:
            模型的预测结果
        """
        return self.pipeline(text)
